import rclpy
from rclpy.node import Node
from sensor_msgs.msg import LaserScan
from sklearn.cluster import DBSCAN
from std_msgs.msg import Float32MultiArray
import numpy as np
import math
from nav_msgs.msg import OccupancyGrid
from geometry_msgs.msg import PointStamped
from tf2_ros import TransformListener, TransformStamped,Buffer
import tf2_geometry_msgs 
import tf2_ros

class RMCostmap(Node):
    def __init__(self):
        super().__init__('rm_costmap')
        self.declare_parameter('expand_range', 10)
        # tf监视器
        self.buffer=Buffer()
        self.tf_listener = TransformListener(self.buffer,self)
        # 创建雷达数据监视器
        self.subscription = self.create_subscription(LaserScan,'/scan',self.scan_callback,50)
        # 创建地图数据监视器
        self.map_subscription = self.create_subscription(OccupancyGrid,'/map',self.map_callback,50)
        # 创建costmap发布方
        self.publisher = self.create_publisher(OccupancyGrid, '/costmap', 50)
        self.map_message = None
        
    # 订阅雷达
    def scan_callback(self, msg):
        if not self.map_message:
            return
        # 新的列表用来存储map
        new_ranges = []
        for i, r in enumerate(msg.ranges):
            # 如果超出雷达范围，设置为'nan'
            if r < msg.range_min or r > msg.range_max or math.isnan(r):
                new_ranges.append(float('nan'))
                continue
            angle = msg.angle_min + i * msg.angle_increment
            # 将极坐标转换为直角坐标系下的位置
            x = r * math.cos(angle)
            y = r * math.sin(angle)
            # 将雷达坐标系下的点转换到地图坐标系(参考坐标系：世界)
            x_world,y_world=self.lidar_to_world(x,y,rclpy.time.Time())
            # 转换参考坐标系，以地图坐标系为观测点
            x_map, y_map = self.world_to_map(x_world,y_world)
            costmap_message=self.map_message
            # 获取雷达上点的序号
            index = y_map * costmap_message.info.width + x_map
            # 膨胀的范围
            expand_range=self.get_parameter('expand_range').get_parameter_value().integer_value
            #膨胀这个点
            for i in range(-expand_range, expand_range, 1):
                # 判断雷达扫描出的点
                if costmap_message.data[i+index] == 0:
                    costmap_message.data[i+index]=100
                for j in range(-expand_range, expand_range, 1):
                    if costmap_message.data[i*costmap_message.info.width+index+j] == 0:
                        costmap_message.data[i*costmap_message.info.width+index+j]=100
        # 创建发布的costmap
        self.publisher.publish(costmap_message)



    def world_to_map(self, x, y):
        wx = x-self.map_message.info.origin.position.x
        wy = y-self.map_message.info.origin.position.y
        mx = int(wx / self.map_message.info.resolution)
        my = int(wy / self.map_message.info.resolution)
        return mx, my
    
    def lidar_to_world(self, x, y,scanTime):
        point_msg=PointStamped()
        point_msg.point.x=x
        point_msg.point.y=y
        try:
            transform = self.buffer.lookup_transform(
                "map", "laser",  # 目标坐标系和源坐标系
                scanTime)
        # 输出转化矩阵
        except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException):
            
            return 0.0,0.0
        # print(transform)
        transform = self.buffer.lookup_transform(
                "map", "laser",  # 目标坐标系和源坐标系
                scanTime)
        point_in_world = tf2_geometry_msgs.do_transform_point(point_msg, transform)
        #point_in_world = self.transform_point(transform, point_msg)
        return point_in_world.point.x,point_in_world.point.y

    def map_callback(self, msg):
        print('map is ok')
        self.map_message = msg

def main(args=None):
    rclpy.init(args=args)
    node = RMCostmap()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()